package cn.denvie.agent.collector;

import cn.denvie.agent.common.AgentConfig;
import cn.denvie.agent.common.AgentContext;
import cn.denvie.agent.common.ByteAgentBuilder;
import cn.denvie.agent.common.MethodBuilder;
import cn.denvie.agent.model.JdbcStatistics;
import javassist.CtClass;
import javassist.CtMethod;

import java.lang.instrument.Instrumentation;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.util.UUID;

/**
 * JDBC数据采集。
 *
 * @author DengZhaoyong
 * @version 1.0.0
 */
public class JdbcCollector extends AbstractByteTransformCollector implements ICollector {

    public static JdbcCollector INSTANCE;

    private final static String[] connection_agent_methods = new String[]{"prepareStatement"};
    private final static String[] prepared_statement_methods = new String[]{"execute", "executeUpdate", "executeQuery"};
    private static final String beginSrc;
    private static final String endSrc;
    private static final String errorSrc;

    private AgentContext context;
    private final String hostIp;
    private final String hostName;

    public JdbcCollector(AgentContext context, Instrumentation instrumentation) {
        super(instrumentation);
        this.context = context;
        this.hostIp = context.getConfig(AgentConfig.HOST_IP, "Undefine");
        this.hostName = context.getConfig(AgentConfig.HOST_NAME, "Undefine");
        INSTANCE = this;
    }

    static {
        beginSrc = "cn.denvie.agent.collector.JdbcCollector inst = cn.denvie.agent.collector.JdbcCollector.INSTANCE;";
        errorSrc = "inst.error(null, e);";
        endSrc = "result = inst.proxyConnection((java.sql.Connection) result);";
    }

    public JdbcStatistics begin(String className, String method) {
        JdbcStatistics jdbcStat = new JdbcStatistics();
        jdbcStat.setTraceId(UUID.randomUUID().toString());
        jdbcStat.setHostIp(hostIp);
        jdbcStat.setHostName(hostName);
        jdbcStat.setBegin(System.currentTimeMillis());
        jdbcStat.setAgentType("jdbc");
        jdbcStat.setRecordTime(System.currentTimeMillis());
        return jdbcStat;
    }

    public void end(JdbcStatistics stat) {
        if (stat == null) {
            return;
        }
        stat.setEnd(System.currentTimeMillis());
        stat.setUseTime(stat.getEnd() - stat.getBegin());
        if (stat.getJdbcUrl() != null) {
            stat.setDatabaseName(getDbName(stat.getJdbcUrl()));
        }
        this.context.submitCollectResult(stat);
    }

    public void error(JdbcStatistics stat, Throwable throwable) {
        if (stat == null) {
            return;
        }
        if (throwable instanceof InvocationTargetException) {
            stat.setError(((InvocationTargetException) throwable).getTargetException().getMessage());
            stat.setErrorType(((InvocationTargetException) throwable).getTargetException().getClass().getName());
        } else {
            stat.setError(throwable.getMessage());
            stat.setErrorType(throwable.getClass().getName());
        }
    }

    public void sendStatistics(JdbcStatistics stat) {
    }

    public Connection proxyConnection(final Connection connection) {
        Object c = Proxy.newProxyInstance(JdbcCollector.class.getClassLoader(),
                new Class[]{Connection.class}, new ConnectionHandler(connection));
        return (Connection) c;
    }

    public PreparedStatement proxyPreparedStatement(final PreparedStatement statement, JdbcStatistics jdbcStat) {
        Object c = Proxy.newProxyInstance(JdbcCollector.class.getClassLoader()
                , new Class[]{PreparedStatement.class}, new PreparedStatementHandler(statement, jdbcStat));
        return (PreparedStatement) c;
    }

    @Override
    public byte[] transform(ClassLoader loader, String className) throws Exception {
        if (!"com.mysql.jdbc.NonRegisteringDriver".equals(className)) {
            return null;
        }
        CtClass ctclass = super.toCtClass(loader, className);
        ByteAgentBuilder byteLoader = new ByteAgentBuilder(className, loader, ctclass);
        CtMethod connectMethod = ctclass.getMethod("connect", "(Ljava/lang/String;Ljava/util/Properties;)Ljava/sql/Connection;");
        // connectMethod.getMethodInfo().getDescriptor();
        MethodBuilder methodBuilder = new MethodBuilder();
        methodBuilder.setBeginSrc(beginSrc);
        methodBuilder.setErrorSrc(errorSrc);
        methodBuilder.setEndSrc(endSrc);
        byteLoader.updateMethod(connectMethod, methodBuilder);
        return byteLoader.toBytecote();
    }

    /**
     * Connection代理处理
     */
    public class ConnectionHandler implements InvocationHandler {
        private final Connection connection;

        private ConnectionHandler(Connection connection) {
            this.connection = connection;
        }

        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            boolean isTargetMethod = false;
            for (String agentm : connection_agent_methods) {
                if (agentm.equals(method.getName())) {
                    isTargetMethod = true;
                }
            }
            Object result = null;
            JdbcStatistics jdbcStat = null;
            try {
                if (isTargetMethod) { // 获取PreparedStatement 开始统计
                    jdbcStat = (JdbcStatistics) JdbcCollector.this.begin(null, null);
                    jdbcStat.setJdbcUrl(connection.getMetaData().getURL());
                    jdbcStat.setSql((String) args[0]);
                }
                result = method.invoke(connection, args);
                // 代理PreparedStatement
                if (isTargetMethod && result instanceof PreparedStatement) {
                    jdbcStat.setPrepared(true);
                    PreparedStatement ps = (PreparedStatement) result;
                    result = proxyPreparedStatement(ps, jdbcStat);
                }
            } catch (Throwable e) {
                JdbcCollector.this.error(jdbcStat, e);
                JdbcCollector.this.end(jdbcStat);
                throw e;
            }
            return result;
        }
    }

    /**
     * PreparedStatement代理处理
     */
    public class PreparedStatementHandler implements InvocationHandler {

        private final PreparedStatement statement;
        private final JdbcStatistics jdbcStat;

        public PreparedStatementHandler(PreparedStatement statement, JdbcStatistics jdbcStat) {
            this.statement = statement;
            this.jdbcStat = jdbcStat;
        }

        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            boolean isTargetMethod = false;
            for (String agentm : prepared_statement_methods) {
                if (agentm.equals(method.getName())) {
                    isTargetMethod = true;
                    break;
                }
            }
            // 收集PreparedStatement执行的参数？
            if ("setString".equals(method.getName())
                    || "setInt".equals(method.getName())
                    || "setLong".equals(method.getName())
                    || "setFloat".equals(method.getName())
                    || "setDouble".equals(method.getName())
                    || "setBigDecimal".equals(method.getName())
                    || "setBoolean".equals(method.getName())
                    || "setShort".equals(method.getName())
                    || "setDate".equals(method.getName())
                    || "setTime".equals(method.getName())
                    || "setTimestamp".equals(method.getName())
                    || "setByte".equals(method.getName())
                    || "setBytes".equals(method.getName())) {
                try {
                    jdbcStat.addParam(new JdbcStatistics.ParamValues((Integer) args[0], args[1]));
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
            // 执行sql语句
            Object result = null;
            try {
                result = method.invoke(statement, args);
            } catch (Throwable e) {
                if (isTargetMethod) {
                    JdbcCollector.this.error(jdbcStat, e);
                }
                throw e;
            } finally {
                if (isTargetMethod) {
                    JdbcCollector.this.end(jdbcStat);
                }
            }
            return result;
        }
    }

    private static String getDbName(String url) {
        int index = url.indexOf("?");
        if (index != -1) {
            url = url.substring(0, index);
        }
        String dbName = url.substring(url.lastIndexOf("/") + 1);
        return dbName;
    }

}